import cv2
import os
import sys
import json
import base64
from tqdm import tqdm
import numpy as np
from skimage import graph, segmentation, color
from skimage.measure import regionprops
import shutil


def rgba_to_hsl(r, g, b, a=255):
    if not (a > 0 and a <= 255):
        return [0]

    """将RGB颜色转换为HSL格式"""
    r_n = r / 255.0
    g_n = g / 255.0
    b_n = b / 255.0
    max_c = max(r_n, g_n, b_n)
    min_c = min(r_n, g_n, b_n)
    delta = max_c - min_c

    # 计算亮度
    l = (max_c + min_c) / 2.0

    if delta == 0:
        h = s = 0
    else:
        # 计算饱和度
        s = delta / (1 - abs(2 * l - 1))
        # 计算色相
        if max_c == r_n:
            h = (g_n - b_n) / delta + (6 if g_n < b_n else 0)
        elif max_c == g_n:
            h = (b_n - r_n) / delta + 2
        else:
            h = (r_n - g_n) / delta + 4
        h /= 6

    # 返回HSL值，保留两位小数以减少文件大小
    return {
        "H": round(h * 360, 2),
        "S": round(s * 100, 2),
        "L": round(l * 100, 2),
        "R": int(r),
        "G": int(g),
        "B": int(b)
    }


def color_segmentation(
    image_path,
    k=8,
    min_pixels=0,
    color_thresh=15
):
    # 读取图像
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    h, w, _ = image.shape

    # 删除并重新创建输出目录
    if os.path.exists('./split'):
        shutil.rmtree('./split')  # 递归删除整个文件夹及其内容
    os.makedirs('./split', exist_ok=True)  # 重新创建空文件夹

    # K-means聚类
    pixel_data = image.reshape((-1, 3)).astype(np.float32)
    criteria = (cv2.TERM_CRITERIA_EPS +
                cv2.TERM_CRITERIA_MAX_ITER, 100, 0.2)
    _, labels, centers = cv2.kmeans(
        pixel_data, k, None, criteria, 10, cv2.KMEANS_RANDOM_CENTERS)
    centers = np.uint8(centers)
    clustered = centers[labels.flatten()].reshape((h, w, 3))
    labels = labels.reshape((h, w)) + 1  # 标签从1开始

    # 构建区域邻接图（RAG）并融合小色块
    if min_pixels > 0:
        rag = graph.rag_mean_color(image, labels)
        regions = regionprops(labels)

        # 标记需要融合的小色块
        small_regions = [r.label for r in regions if r.area < min_pixels]
        neighbor_map = {r.label: [] for r in regions}

        # 查找每个小色块的相邻色块
        for region in regions:
            if region.label in small_regions:
                boundary_labels = np.unique(labels[max(0, region.bbox[0]-1):min(h, region.bbox[2]+1),
                                                   max(0, region.bbox[1]-1):min(w, region.bbox[3]+1)])
                neighbor_map[region.label] = [
                    x for x in boundary_labels if x != region.label and x not in small_regions]

        # 融合到最相似的相邻色块
        for small_label in small_regions:
            if not neighbor_map[small_label]:
                continue
            # 计算颜色相似度
            small_color = rag.nodes[small_label]['mean color']
            min_diff = float('inf')
            best_neighbor = None
            for neighbor in neighbor_map[small_label]:
                neighbor_color = rag.nodes[neighbor]['mean color']
                diff = np.linalg.norm(small_color - neighbor_color)
                if diff < min_diff and diff < color_thresh:
                    min_diff = diff
                    best_neighbor = neighbor
            # 执行融合
            if best_neighbor:
                labels[labels == small_label] = best_neighbor

    height, width = image.shape[:2]

    # 保存色块（保持原分辨率）
    output_segments = {
    }
    for label in tqdm(np.unique(labels)):
        hsl_img = np.zeros((height, width, 1), dtype=np.float32)

        if label == 0:
            continue  # 跳过背景
        mask = np.where(labels == label, 255, 0).astype(np.uint8)
        mean_color = cv2.mean(image, mask)[:3]

        # 创建带透明通道的色块
        segmented_img = np.zeros((h, w, 4), dtype=np.uint8)
        segmented_img[..., :3] = mean_color
        segmented_img[..., 3] = mask

        json_filename = f'segment_{label:03d}'

        # 保存图片
        # output_path = f'./split/segment_{json_filename}.png'
        # cv2.imwrite(output_path, cv2.cvtColor(
        #     segmented_img, cv2.COLOR_RGBA2BGRA))

        # 保存json
        # output_json_path = f'./split/segment_{json_filename}.json'

        # RGBA json
        # with open(output_json_path, 'w') as f:
        #     json.dump(segmented_img.tolist(), f)

        # HSL json
        split_color = None
        for i in range(height):
            for j in range(width):
                r, g, b, a = segmented_img[i, j, :4]

                if a == 0:
                    hsl_img[i, j] = [0]
                else:
                    if split_color is None:
                        split_color = rgba_to_hsl(r, g, b, a)
                    hsl_img[i, j] = [1]

                # hsl = rgba_to_hsl(r, g, b, a)
                # hsl_img[i, j] = hsl
        # output_segments[json_filename] = hsl_img.tolist()

        # 1. 将图像转换为PNG字节流
        # _, buffer = cv2.imencode('.png', cv2.cvtColor(
        #     segmented_img, cv2.COLOR_RGBA2BGRA))
        # # 2. 编码为base64字符串
        # img_base64 = base64.b64encode(buffer).decode('utf-8')

        output_segments[json_filename] = {
            "offset": [0, 0, 0],
            "split_color": split_color,
            "data": hsl_img.tolist()
        }

        # with open(output_json_path, 'w') as f:
        #     json.dump(hsl_img.tolist(), f)

    # with open('./split/ret.json', 'w') as f:
    #     json.dump(output_segments, f)

    print(json.dumps(output_segments))
    return output_segments


# 使用示例
if __name__ == "__main__":
    conf = json.loads(sys.argv[1])

    # 示例1：K-means聚类+小色块融合
    color_segmentation(
        conf['path'],
        conf['k'],
        conf['min_px'],
        conf['color_thresh']
    )
